import threading
import queue
import time # 新增导入

from config import SERVER_IPS
from logger_setup import logger
from utils import ServerTaskFailedError
from task_processor import process_notification_background, send_wechat_notification, send_status_notification # 导入 send_status_notification

# 任务队列和服务器队列
task_queue = queue.Queue()  # 任务队列，存储待处理任务
server_queue = queue.Queue()  # 服务器队列，管理可用服务器资源
processed_task_ids = set() # 用于存储已处理的任务ID (这仍然可以用于某些场景，例如防止已完成任务被错误地重新激活)
current_task_ids_in_queue = set() # 新增：用于跟踪当前在任务队列中的任务ID
task_id_lock = threading.Lock() # 新增：用于同步访问 current_task_ids_in_queue

def add_task_to_queue(task_data):
    """
    将任务添加到任务队列，并记录其ID。
    """
    task_id = str(task_data.get("id", "未知ID"))
    with task_id_lock:
        if task_id in current_task_ids_in_queue:
            logger.info(f"[任务添加] 检测到任务池中已存在任务ID: {task_id} ({task_data.get('name', '未知名称')})，将忽略此重复任务的添加。")
            return False # 表示任务未添加
        task_queue.put(task_data)
        current_task_ids_in_queue.add(task_id)
        logger.info(f"[任务添加] 任务 ID {task_id} ({task_data.get('name', '未知名称')}) 已成功加入任务队列和ID跟踪集合。")
        return True # 表示任务已添加

def process_notification_background_wrapper(task_data, server_ip):
    """
    任务处理包装器，负责错误处理与服务器资源管理
    
    Args:
        task_data: 任务数据字典
        server_ip: 分配用于处理任务的服务器IP
    """
    try:
        process_notification_background(task_data, server_ip)
    except ServerTaskFailedError as e_server_fail:
        task_id = task_data.get("id", "未知ID")
        task_name = task_data.get("name", "未知名称")
        logger.warning(f"[任务包装器] 任务 ID {task_id} ({task_name}) 在服务器 {server_ip} 上处理失败，将尝试重新排队。原始错误: {e_server_fail.original_exception}")

        attempted_servers = task_data.get("_attempted_servers", [])
        if server_ip not in attempted_servers:
            attempted_servers.append(server_ip)
        task_data["_attempted_servers"] = attempted_servers
        task_id_str = str(task_id)

        if len(attempted_servers) < len(SERVER_IPS):
            task_id_str = str(task_id) # 确保 task_id 类型一致性
            # 如果任务重新排队，它不应存在于 processed_task_ids 中。
            # 它会通过 add_task_to_queue 重新加入 current_task_ids_in_queue
            if task_id_str in processed_task_ids:
                try:
                    processed_task_ids.remove(task_id_str)
                    logger.info(f"[任务包装器] 任务 ID {task_id_str} ({task_name}) 从 processed_task_ids 中移除，因为它将被重新排队。")
                except KeyError:
                    logger.warning(f"[任务包装器] 尝试从 processed_task_ids 移除任务 ID {task_id_str} 时发生 KeyError，可能已被其他逻辑移除。")
            # task_queue.put(task_data) # 改为通过 add_task_to_queue 添加
            if add_task_to_queue(task_data): # 如果重新排队成功
                 logger.info(f"[任务包装器] 任务 ID {task_id} ({task_name}) 已重新加入队列。已尝试服务器: {attempted_servers}")
            else:
                # 如果因为 current_task_ids_in_queue 中已存在而无法重新排队 (理论上不应发生，除非逻辑冲突)
                logger.error(f"[任务包装器] 尝试重新排队任务 ID {task_id} ({task_name}) 失败，因为它已存在于 current_task_ids_in_queue。这可能是一个逻辑错误。")

        else:
            logger.error(f"[任务包装器] 任务 ID {task_id} ({task_name}) 已在所有服务器 ({attempted_servers}) 上尝试失败，不再重新排队。最终错误来自 {server_ip}: {e_server_fail.original_exception}")
            send_wechat_notification(task_id, task_name, 'fail',
                                     f'任务在所有服务器上均处理失败 (最后尝试: {server_ip})',
                                     error_reason=f"已尝试: {attempted_servers}. 末次错误: {str(e_server_fail.original_exception)}")
            # 当所有服务器都失败后，从 processed_task_ids 中移除，允许任务ID被重新提交
            if task_id_str in processed_task_ids:
                try:
                    processed_task_ids.remove(task_id_str)
                    logger.info(f"[任务包装器] 任务 ID {task_id_str} ({task_name}) 因所有服务器尝试失败，已从 processed_task_ids 中移除。")
                except KeyError:
                    logger.warning(f"[任务包装器] 尝试从 processed_task_ids 移除最终失败的任务 ID {task_id_str} 时发生 KeyError。")

    except Exception as e_global_wrapper:
        task_id = task_data.get("id", "未知ID")
        task_name = task_data.get("name", "未知名称")
        task_id_str = str(task_id) # 确保 task_id 类型一致性
        logger.error(f"[任务包装器] 处理任务 ID {task_id} ({task_name}) (服务器 {server_ip}) 时发生未捕获的意外错误: {e_global_wrapper}", exc_info=True)
        # 发生未捕获的意外错误时，也从 processed_task_ids 中移除，允许任务ID被重新提交
        if task_id_str in processed_task_ids:
            try:
                processed_task_ids.remove(task_id_str)
                logger.info(f"[任务包装器] 任务 ID {task_id_str} ({task_name}) 因未捕获的意外错误，已从 processed_task_ids 中移除。")
            except KeyError:
                logger.warning(f"[任务包装器] 尝试从 processed_task_ids 移除因意外错误失败的任务 ID {task_id_str} 时发生 KeyError。")

    finally:
        logger.info(f"[服务器管理] 服务器 {server_ip} 完成尝试任务 ID {task_data.get('id', '未知ID')}，已返回服务器队列。")
        server_queue.put(server_ip)

def dispatch_tasks():
    """
    任务调度主线程，管理任务队列和服务器资源分配。
    一旦任务从task_queue中取出，它将持续等待，直到找到一个合适的（未尝试过的）服务器。
    """
    for ip in SERVER_IPS:
        server_queue.put(ip)
    
    logger.info(f"任务调度器已启动，管理服务器: {SERVER_IPS}")

    while True:
        current_task_data = None
        assigned_server = None
        # 标记是否已为当前任务在本次调度等待中发送过"排队中"的通知和相关日志
        waiting_notification_and_log_sent_this_cycle = False
        
        try:
            current_task_data = task_queue.get() 
            task_id = current_task_data.get("id", "未知ID")
            task_name = current_task_data.get("name", "未知名称")
            task_id_str = str(task_id) # 统一使用字符串类型的task_id

            # 任务已从队列取出，准备处理，将其从 current_task_ids_in_queue 中移除
            with task_id_lock:
                if task_id_str in current_task_ids_in_queue:
                    current_task_ids_in_queue.remove(task_id_str)
                    logger.info(f"[任务调度] 任务 ID {task_id_str} ({task_name}) 已从 current_task_ids_in_queue 中移除，准备处理。")
                else:
                    # 这可能表示任务在加入队列和取出队列之间被异常移除了，或者从未正确加入
                    logger.warning(f"[任务调度] 任务 ID {task_id_str} ({task_name}) 在从队列取出时，未在 current_task_ids_in_queue 中找到。")

            if task_id_str not in processed_task_ids: # 仅在任务ID确实是新的（未曾处理完成或失败）时添加
                 processed_task_ids.add(task_id_str)
                 logger.info(f"[任务调度] 新任务ID: {task_id_str} ({task_name}) 已记录到 processed_task_ids，准备处理。")
            else:
                 logger.info(f"[任务调度] 任务ID: {task_id_str} ({task_name}) 已存在于 processed_task_ids 中，但由于任务池中无重复，仍将处理。")

            logger.info(f"[任务调度] 任务 ID {task_id_str} ({task_name}) 已从主队列获取，开始为其寻找合适服务器。")

            while True:
                attempted_by_current_task = current_task_data.get("_attempted_servers", [])

                if len(SERVER_IPS) > 0 and len(attempted_by_current_task) >= len(SERVER_IPS):
                    logger.error(f"[任务调度] 任务 ID {task_id_str} ({task_name}) 已尝试过所有服务器 {attempted_by_current_task}。此任务不应被重排但仍存在于队列中。将标记为完成并丢弃此异常任务。")
                    task_queue.task_done() # 标记此任务已处理（通过丢弃）
                    # 如果任务被丢弃，也应从 processed_task_ids 中移除，因为它实际上未被处理
                    # 并且也应确保它不在 current_task_ids_in_queue 中 （理论上此时已经移除了）
                    with task_id_lock:
                        if task_id_str in current_task_ids_in_queue: # 双重检查
                            current_task_ids_in_queue.remove(task_id_str)
                            logger.info(f"[任务调度] 丢弃任务 {task_id_str} 时，从 current_task_ids_in_queue 中移除。")
                        if task_id_str in processed_task_ids:
                            try:
                                processed_task_ids.remove(task_id_str)
                                logger.info(f"[任务调度] 丢弃任务 {task_id_str} 时，从 processed_task_ids 中移除。")
                            except KeyError:
                                pass # 可能已经被其他逻辑移除
                    current_task_data = None # 清理，以防影响外部循环的判断
                    break # 跳出服务器分配循环，由于 current_task_data 为 None，外部 if 会阻止处理

                candidate_server = None
                try:
                    candidate_server = server_queue.get(block=True, timeout=None) 
                    
                    if candidate_server not in attempted_by_current_task:
                        assigned_server = candidate_server 
                        logger.info(f"[任务调度] 任务 ID {task_id_str} ({task_name}) 分配给可用且合适的服务器 {assigned_server}。")
                        break 
                    else:
                        server_queue.put(candidate_server) # 不合适的服务器，立即归还
                        candidate_server = None 
                        
                        if not waiting_notification_and_log_sent_this_cycle: # 只记录日志，不发送通知
                            logger.info(f"[任务调度] 任务 ID {task_id_str} ({task_name}) 正在等待一个之前未尝试过的空闲服务器。当前可用服务器均已尝试过。")
                            waiting_notification_and_log_sent_this_cycle = True
                        time.sleep(0.2) 

                except queue.Empty:
                    logger.warning(f"[任务调度] 任务 ID {task_id_str} ({task_name}) 等待服务器时 server_queue.get() 意外返回 Empty。将重试等待。")
                    time.sleep(1) 
                except Exception as e_get_server:
                    logger.error(f"[任务调度] 为任务 ID {task_id_str} ({task_name}) 获取服务器过程中发生错误: {e_get_server}", exc_info=True)
                    if candidate_server: 
                        server_queue.put(candidate_server)
                    # task_queue.put(current_task_data) # 如果获取服务器失败，任务需要重新加入队列
                    # current_task_data = None # 清理
                    # 改为通过 add_task_to_queue 重新添加，以维护 current_task_ids_in_queue
                    if current_task_data: # 确保 current_task_data 不是 None
                        logger.info(f"[任务调度] 为任务 ID {task_id_str} ({task_name}) 获取服务器失败，将尝试重新加入任务队列。")
                        if not add_task_to_queue(current_task_data):
                            # 如果因重复或其他原因无法重新添加，则记录错误
                            logger.error(f"[任务调度] 尝试将获取服务器失败的任务 ID {task_id_str} ({task_name}) 重新加入队列失败。该任务可能丢失。")
                        current_task_data = None # 清理
                    time.sleep(2)
                    break 

            if current_task_data and assigned_server:
                logger.info(f"[任务调度] 正式为任务 ID {task_id_str} ({task_name}) 在服务器 {assigned_server} 上启动处理线程。")
                thread = threading.Thread(target=process_notification_background_wrapper, args=(current_task_data.copy(), assigned_server))
                thread.start()
                task_queue.task_done() # 标记任务完成
                current_task_data = None # 确保处理完后清空，避免影响下一次循环的判断
                assigned_server = None   

        except Exception as e_dispatch_loop:
            logger.error(f"[任务调度] 调度主循环发生严重错误: {e_dispatch_loop}", exc_info=True)
            if current_task_data: 
                logger.error(f"因调度主循环错误，尝试将任务 {current_task_data.get('id', '未知ID')} 放回主队列。")
                # task_queue.put(current_task_data) # 改为通过 add_task_to_queue
                if not add_task_to_queue(current_task_data):
                    logger.error(f"[任务调度] 因调度主循环错误，尝试将任务 {current_task_data.get('id', '未知ID')} 重新加入队列失败。该任务可能丢失。")

            time.sleep(5)